# -*- coding: utf-8 -*-
import numpy as np
import scipy.io as sio
from .CSPTrain import csp_train


def csp_spatial_filter(data_x, csp_proj_matrix, filter_bank=False):
    """
    返回 processor 空间滤波后的数据

    输入参数
    ----------
    data_x: T×N×L ndarray(或单个trial T×N)
           T: 采样点数  N: 通道数  L: 训练数据trial总数
    csp_proj_matrix: 2m×N
           processor 投影矩阵
    返回值
    ----------
    x_after_csp: L×2m (或单个trial 1D 2m)  空间滤波后的数据
    """
    feature_len = csp_proj_matrix.shape[0]
    if filter_bank:
        filter_num = data_x.shape[2]
        if len(data_x.shape) == 4:
            trial_size = data_x.shape[3]
            x_after_csp = np.zeros([trial_size, filter_num, feature_len])
            for k in range(filter_num):
                for i in range(trial_size):
                    wk = csp_proj_matrix[:, :, k]  # 投影矩阵 2m×channel
                    eki = data_x[:, :, k, i]  # EEG信号 sample×channel
                    x_after_csp[i, k, :] = get_feature(wk, eki)
        else:
            x_after_csp = np.zeros([feature_len, filter_num])
            for k in range(filter_num):
                wk = csp_proj_matrix[:, :, k]  # 投影矩阵
                eki = data_x[:, :, k]  # EEG信号
                x_after_csp[k, :] = get_feature(wk, eki)
    else:
        if len(data_x.shape) == 3:
            trial_size = data_x.shape[2]
            x_after_csp = np.zeros([trial_size, feature_len])
            for i in range(trial_size):
                x_after_csp[i, :] = get_feature(csp_proj_matrix, data_x[:, :, i])
        else:
            feature = get_feature(csp_proj_matrix, data_x)
            x_after_csp = np.zeros([1, feature_len])
            x_after_csp[0, :] = feature
    return x_after_csp


def get_feature(wk, eki):
    wk_t = np.transpose(wk)  # 转置
    eki_t = np.transpose(eki)
    z1 = np.dot(np.dot(np.dot(wk, eki_t), eki), wk_t)
    z2 = np.log(np.diag(z1) / np.trace(z1))
    return z2


if __name__ == '__main__':
    trainx = sio.loadmat(r'D:\Myfiles\WorkSpace\Codes\PythonProjects\Data\trainx.mat')
    trainy = sio.loadmat(r'D:\Myfiles\WorkSpace\Codes\PythonProjects\Data\trainy.mat')
    train_x = trainx['train_x']  # shape(750,22,60)
    train_y = trainy['train_y']  # shape(1,60)
    m = 3  # processor 参数
    csp_ProjMatrix = csp_train(train_x, train_y.ravel(), m)
    xAfterCSP = csp_spatial_filter(train_x, csp_ProjMatrix)
    print(xAfterCSP.shape())